PiPNN 3/6: add core graph construction - #1290
Conversation
There was a problem hiding this comment.
Pull request overview
This PR introduces the “core” PiPNN build pipeline in the diskann-pipnn crate, wiring together deterministic partitioning, leaf-local candidate construction, and final pruning into a public build_graph API with a validated build context.
Changes:
- Adds
PiPNNConfigvalidation and aPiPNNBuildContextthat binds PiPNN policy to DiskANN graph pruning policy and a caller-owned Rayon thread pool. - Implements the three main stages: partitioning (
partitioning.rs), leaf candidate construction (leaf_build.rs), and final pruning via shared Vamana robust prune (finalization.rs). - Adds comprehensive unit/integration tests and a Criterion benchmark for core scenarios; updates dependencies, lockfile, and mutation-test exclusions.
Reviewed changes
Copilot reviewed 13 out of 14 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| diskann-pipnn/src/lib.rs | Adds public PiPNN API (PiPNNConfig, PiPNNBuildContext, build_graph) and stage orchestration. |
| diskann-pipnn/src/partitioning.rs | Implements deterministic overlapping partition construction and leader assignment/scatter. |
| diskann-pipnn/src/partitioning/tests.rs | Adds unit tests covering partition determinism, invariants, error cases, and helpers. |
| diskann-pipnn/src/leaf_build.rs | Builds leaf-local symmetric k-NN candidates and accumulates global candidates safely in parallel. |
| diskann-pipnn/src/leaf_build/tests.rs | Adds unit tests for candidate correctness, invariants, type support, and error handling. |
| diskann-pipnn/src/finalization.rs | Orders/prunes candidate rows using shared robust_prune and validates candidate IDs/shape. |
| diskann-pipnn/src/finalization/tests.rs | Adds unit tests for pruning behavior and candidate validation failures. |
| diskann-pipnn/src/tests.rs | Tests effective_metric behavior for integer cosine-normalized handling. |
| diskann-pipnn/tests/config.rs | Integration tests for config validation and graph-policy compatibility checks. |
| diskann-pipnn/tests/build_graph.rs | Integration tests for end-to-end graph building, invariants, determinism, and type/metric support. |
| diskann-pipnn/benches/core.rs | Adds a Criterion benchmark for stage-focused core build scenarios. |
| diskann-pipnn/Cargo.toml | Updates crate dependencies/dev-dependencies and registers the new core benchmark target. |
| Cargo.lock | Records dependency graph changes for the updated diskann-pipnn crate dependencies. |
| .cargo/mutants.toml | Adds mutation-test exclusions for key PiPNN public boundary checks and partitioning invariants. |
Comments suppressed due to low confidence (1)
diskann-pipnn/src/partitioning.rs:604
size_of::<f32>()is used without being in scope (nouse std::mem::size_of;and not qualified), so this function won’t compile as written.
fn assignment_stripe_rows(leaders: usize) -> usize {
(ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::<f32>()))
.clamp(MIN_ASSIGNMENT_STRIPE_ROWS, MAX_ASSIGNMENT_STRIPE_ROWS)
}
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| *scale = FastL2NormSquared.evaluate(row); | ||
| if metric == Metric::Cosine { | ||
| *scale = scale.sqrt(); | ||
| } | ||
| } |
5be0c9c to
50047c6
Compare
| fn assignment_stripe_rows(leaders: usize) -> usize { | ||
| (ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::<f32>())) | ||
| .clamp(MIN_ASSIGNMENT_STRIPE_ROWS, MAX_ASSIGNMENT_STRIPE_ROWS) | ||
| } |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 14 out of 15 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (2)
diskann-pipnn/src/partitioning.rs:637
size_ofis used without being in scope (std::mem::size_of), which will not compile. Qualify the call or import it.
fn assignment_stripe_rows(leaders: usize) -> usize {
(ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::<f32>()))
.clamp(MIN_ASSIGNMENT_STRIPE_ROWS, MAX_ASSIGNMENT_STRIPE_ROWS)
}
diskann-pipnn/src/partitioning.rs:19
Normis imported but never used in this module, which will tripunused_importswarnings (and can become CI failures under-D warnings). Remove it from the import list.
use diskann::{utils::VectorRepr, ANNError, ANNResult};
use diskann_linalg::Transpose;
use diskann_utils::views::MatrixView;
use diskann_vector::{distance::Metric, norm::FastL2NormSquared, Norm};
use rand::{prelude::IndexedRandom, SeedableRng};
use rayon::prelude::*;
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 14 out of 15 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (1)
diskann-pipnn/src/partitioning.rs:390
gather_rowsusesTypeId::of::<T>(), which implicitly requiresT: 'static. Making that bound explicit here avoids surprising/indirect trait-bound errors later and matches the publicbuild_graphboundary (which already requires'static).
fn gather_rows<T>(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()>
where
T: VectorRepr,
1324668 to
857e200
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 14 out of 15 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (1)
diskann-pipnn/src/partitioning.rs:641
size_of::<f32>()is used without being imported or qualified, which will fail to compile. Qualify it withstd::mem::size_of(or add an explicit import).
let rows = ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::<f32>());
| use diskann::{utils::VectorRepr, ANNError, ANNResult}; | ||
| use diskann_linalg::Transpose; | ||
| use diskann_utils::views::MatrixView; | ||
| use diskann_vector::{distance::Metric, norm::FastL2NormSquared, Norm}; |
857e200 to
f642204
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 14 out of 15 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (4)
diskann-pipnn/src/leaf_build.rs:222
build_leafis executed from a Rayon parallel context (viabuild_leaf_candidates), so it should also explicitly requireT: Send + Syncto reflect the actual thread-safety requirement.
where
T: VectorRepr + 'static,
{
diskann-pipnn/src/leaf_build/tests.rs:154
assert_source_typeforwardsTinto the parallel leaf build path, so it should also includeSend + Syncbounds to match the production requirements.
fn assert_source_type<T>(data: &[T])
where
T: diskann::utils::VectorRepr + 'static,
{
diskann-pipnn/src/leaf_build.rs:193
build_leaf_candidatesuses Rayon parallel iteration overdata, soTmust beSend + Sync. Making this explicit in the signature avoids confusing trait-bound errors at call sites and documents the thread-safety requirement.
This issue also appears on line 220 of the same file.
where
T: VectorRepr + 'static,
{
diskann-pipnn/src/leaf_build/tests.rs:35
- This test helper calls
build_leaf_candidates, which (via Rayon) requiresT: Send + Sync. Add the bounds here so the test continues to compile once the production signature is tightened.
This issue also appears on line 151 of the same file.
where
T: diskann::utils::VectorRepr + 'static,
{
| fanout: usize, | ||
| leaders: usize, | ||
| ) -> ANNResult<Vec<Vec<u32>>> { | ||
| let mut sizes = filled_vec(leaders, 0usize)?; |
f642204 to
b1181d6
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 14 out of 15 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (1)
diskann-pipnn/src/finalization.rs:76
prune::robust_prunerequires the candidate pool to be sorted by increasing distance, butpoolis currently populated in candidate-ID order (from theAdjacencyList) and never sorted by the computed distances. This can change pruning behavior substantially and break determinism/quality.
let pool = workspace.prune.candidates_mut();
pool.clear();
pool.try_reserve(row.len()).map_err(ANNError::opaque)?;
pool.extend(row.iter().copied().map(|candidate| {
Neighbor::new(
candidate,
distance.evaluate_similarity(source_vector, data.row(candidate as usize)),
)
}));
let candidate_count = pool.len();
let mut context = workspace.prune.as_context(candidate_count);
prune::robust_prune(
| fn assignment_stripe_rows(leaders: usize) -> usize { | ||
| let rows = ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::<f32>()); | ||
| let rows = if rows.is_power_of_two() { | ||
| rows | ||
| } else { | ||
| rows.next_power_of_two() / 2 | ||
| }; | ||
| rows.clamp(MIN_ASSIGNMENT_STRIPE_ROWS, MAX_ASSIGNMENT_STRIPE_ROWS) | ||
| } |
| /// Number of nearest leaders retained at each overlapping partition level. | ||
| pub fanout: Vec<usize>, | ||
| /// Number of nearest neighbors selected within each leaf. | ||
| pub k: usize, |
There was a problem hiding this comment.
Request to use something like partition_k to not overload k
| /// Fraction of a cluster sampled as partition leaders. | ||
| pub p_samp: f64, | ||
| /// Number of nearest leaders retained at each overlapping partition level. | ||
| pub fanout: Vec<usize>, |
There was a problem hiding this comment.
Could you explain why this is a vector? What would it mean if I passed in an arbitrary vector here?
| self.c_min, self.c_max | ||
| ))); | ||
| } | ||
| if !self.p_samp.is_finite() || !(0.0..=1.0).contains(&self.p_samp) || self.p_samp == 0.0 { |
There was a problem hiding this comment.
Isn't the middle statement sufficient? Why do we need to explicitly check finite and nonzero?
| pool: &'a ThreadPool, | ||
| ) -> ANNResult<Self> { | ||
| config.validate()?; | ||
| if !graph.alpha().is_finite() || graph.alpha() < 1.0 { |
There was a problem hiding this comment.
0.0 < Alpha < 1.0 should not be disallowed. It is meaningful. Also, the Config should check for these issues already. If it doesn't, the errors should be moved into the validation of the config itself.
| /// Policy owned by the partition stage. Leaf-neighbor and merge settings do | ||
| /// not cross this boundary. | ||
| #[derive(Clone, Debug)] | ||
| pub(crate) struct PartitionConfig { | ||
| c_max: usize, | ||
| c_min: usize, | ||
| p_samp: f64, | ||
| fanout: Vec<usize>, | ||
| replicas: usize, | ||
| } | ||
|
|
||
| impl From<&PiPNNConfig> for PartitionConfig { | ||
| fn from(config: &PiPNNConfig) -> Self { | ||
| Self { | ||
| c_max: config.c_max, | ||
| c_min: config.c_min, | ||
| p_samp: config.p_samp, | ||
| fanout: config.fanout.clone(), | ||
| replicas: config.replicas, | ||
| } | ||
| } |
There was a problem hiding this comment.
Why can't we just use PiPNNConfig? What would go wrong if we did?
17723cf to
a7ce31c
Compare
|
@SeliMeli please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
Contributor License AgreementContribution License AgreementThis Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
|
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 13 out of 14 changed files in this pull request and generated no new comments.
Suppressed comments (1)
diskann-pipnn/src/partitioning.rs:172
- The doc comment claims coverage is validated ("every input point must remain covered once per replica"), but
partition()currently only callsvalidate_leaves(&leaves, config.c_max)which checks empties/oversized leaves. Either add an explicit per-replica coverage check, or adjust the docs so they don’t promise validation that isn’t performed.
/// Levels beyond `fanout.len()` retain one leader assignment. Completed small
/// leaves are merged without exceeding `c_max`; every input point must remain
/// covered once per replica. The caller installs the operation in its pool.
Adds the provider-independent PiPNN graph-construction pipeline.
Code map
lib.rsdefines configuration, validates graph-policy compatibility, and orchestratespartition → leaf candidates → finalizationinside the caller-owned Rayon pool.partitioning.rsruns deterministic replicas. Each oversized work item samples leaders, gathers point/leader rows, uses GEMM pluspartition_kernelfor assignments, and recurses only on oversized clusters.global_merge_smallcombines sub-c_minleaves without exceedingc_max; final validation rejects empty/oversized leaves.leaf_build.rsgives each Rayon job reusable buffers. Active prefixes are passed explicitly because numeric buffers retain their high-water length. Each leaf gathers rows, computes lowerA · Aᵀ, runs the dual-endpoint top-k kernel, and merges symmetric candidates.finalization.rsleaves degree-bounded rows unchanged and sends only overfull rows through the shared RobustPrune kernel.Review path
leaf_build, check the sorted-ID duplicate fast path and its HashSet fallback, plus every active-prefix slice after grow-only buffer reuse.MapInitstate, and the outer leaf value is consumed by the stage. There is no TLS or cleanup broadcast.Stack 3/6: #1288 → #1291